﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using System.Diagnostics;
using System.Data.Cassandra.Connection;
using Thrift.Transport;
using Thrift.Protocol;
using System.Data.Cassandra.Extensions;
using System.Data.Cassandra.Connection.Nodes;

namespace System.Data.Cassandra.Connection.Pool
{
    public class ConnectionPool : IDisposable
    {
        [DebuggerDisplay("{Host}({IdleClients.Count}/{Clients.Count})")]
        private class HostClients
        {
            private readonly ConnectionPool connectionPool;
            private readonly List<Apache.Cassandra.Cassandra.Client> clients;

            public readonly string Host;
            public readonly Queue<Apache.Cassandra.Cassandra.Client> IdleClients;

            private Exception lastError;
            private DateTime lastErrorAtUtc;

            public ConnectionPool ConnectionPool
            {
                get
                {
                    return connectionPool;
                }
            }
            public int Count
            {
                get
                {
                    return clients.Count;
                }
            }
            public bool IsLastErrorRecent
            {
                get
                {
                    DateTime errorAtUtc = lastErrorAtUtc;
                    if (errorAtUtc > DateTime.MinValue)
                    {
                        TimeSpan t = DateTime.UtcNow.Subtract(errorAtUtc);
                        if (t < connectionPool.disableConnectionAfterError)
                        {
                            return true;
                        }
                        this.lastErrorAtUtc = DateTime.MinValue;
                    }
                    return false;
                }
            }

            public HostClients(ConnectionPool pool, string host)
            {
                connectionPool = pool;
                Host = host;
                int maxPoolCountPerHost = pool.MaxPoolCountPerHost;
                int capacity = Math.Min(32, maxPoolCountPerHost);
                clients = new List<Apache.Cassandra.Cassandra.Client>(capacity);
                IdleClients = new Queue<Apache.Cassandra.Cassandra.Client>(capacity);
            }

            public override string ToString()
            {
                return string.Format("{0}({1}/{2})", this.Host, this.IdleClients.Count, this.clients.Count);
            }

            private Apache.Cassandra.Cassandra.Client Add()
            {
                Stopwatch sw = Stopwatch.StartNew();
                CassandraConnectionConfig connectionConfig = connectionPool.ConnectionConfig;
                TTransport tTransport = new TFramedTransport(new TSocket(Host, connectionConfig.Port));
                TBinaryProtocol prot = new TBinaryProtocol(tTransport);
                Apache.Cassandra.Cassandra.Client client = new Apache.Cassandra.Cassandra.Client(prot);
                try
                {
                    tTransport.Open();
                    client.set_keyspace(connectionConfig.Keyspace);
                }
                catch (Exception err)
                {
                    OnConnectFailed(this, sw, err);
                    MarkError(err);
                    client.Close();
                    throw;
                }
                OnConnected(this, sw);
                clients.Add(client);
                return client;
            }

            public Apache.Cassandra.Cassandra.Client Borrow(int maxConnections, bool fIgnoreRecentError)
            {
                Queue<Apache.Cassandra.Cassandra.Client> idleClients = IdleClients;
                Apache.Cassandra.Cassandra.Client result;
                if (idleClients.Count > 0)
                {
                    Apache.Cassandra.Cassandra.Client client = idleClients.Dequeue();
                    try
                    {
                        TTransport transport = client.InputProtocol.Transport;
                        if (!transport.IsOpen)
                        {
                            transport.Open();
                        }
                    }
                    catch (Exception ex)
                    {
                        MarkError(ex);
                        client.Close();
                        throw;
                    }
                    result = client;
                }
                else
                {
                    if (!fIgnoreRecentError && IsLastErrorRecent)
                    {
                        result = null;
                    }
                    else
                    {
                        if (maxConnections > 0 && clients.Count >= maxConnections)
                        {
                            result = null;
                        }
                        else
                        {
                            result = Add();
                        }
                    }
                }
                return result;
            }

            public void MarkError(Exception err)
            {
                lastError = err;
                lastErrorAtUtc = DateTime.UtcNow;
            }

            public void CloseAllIdleClients()
            {
                Queue<Apache.Cassandra.Cassandra.Client> idleClients = this.IdleClients;
                foreach (Apache.Cassandra.Cassandra.Client current in idleClients)
                {
                    current.Close();
                }
                idleClients.Clear();
            }

            public void CloseAllClients()
            {
                List<Apache.Cassandra.Cassandra.Client> list = this.clients;
                foreach (Apache.Cassandra.Cassandra.Client current in list)
                {
                    current.Close();
                }
                list.Clear();
            }
        }

        private readonly object borrowReturnLock = new object();
        private readonly CassandraConnectionConfig connectionConfig;
        private readonly ICassandraNodeEnumerator nodeEnumerator;

        private AutoResetEvent autoResetEvent = new AutoResetEvent(false);
        private TimeSpan disableConnectionAfterError = TimeSpan.FromMinutes(5.0);
        private Dictionary<string, HostClients> clientsByHost = new Dictionary<string, HostClients>();
        private bool disposed;

        public static event EventHandler<CassandraConnectionEventArgs> Connected;
        public static event EventHandler<CassandraConnectionEventArgs> ConnectionFailed;

        public CassandraConnectionConfig ConnectionConfig
        {
            get
            {
                return connectionConfig;
            }
        }
        public string FirstHost
        {
            get
            {
                return connectionConfig.Hosts.First();
            }
        }
        public int FirstPort
        {
            get
            {
                return connectionConfig.Port;
            }
        }
        public string Keyspace
        {
            get
            {
                return connectionConfig.Keyspace;
            }
        }
        public TimeSpan Timeout
        {
            get
            {
                return this.connectionConfig.Timeout;
            }
        }
        public int TimeoutMillisecond
        {
            get
            {
                return (int)this.connectionConfig.Timeout.TotalMilliseconds;
            }
        }
        public int MinPoolCountPerHost
        {
            get
            {
                return this.connectionConfig.MinPoolCountPerHost;
            }
        }
        public int MaxPoolCountPerHost
        {
            get
            {
                return this.connectionConfig.MaxPoolCountPerHost;
            }
        }
        public TimeSpan DisableConnectionAfterError
        {
            get
            {
                return this.disableConnectionAfterError;
            }
            set
            {
                this.disableConnectionAfterError = value;
            }
        }
        public ICollection<string> Nodes
        {
            get
            {
                return this.clientsByHost.Keys;
            }
        }
        public ICassandraNodeEnumerator NodeEnumerator
        {
            get
            {
                return this.nodeEnumerator;
            }
        }

        internal bool IsDisposed
        {
            get
            {
                return this.disposed;
            }
        }

        public ConnectionPool(CassandraConnectionConfig config)
        {
            connectionConfig = config;
            autoResetEvent.Set();
            UpdateNodes(config.Hosts);
            nodeEnumerator = config.CreateNodeEnumerator(this);
            TryUpdateNodes();
        }

        ~ConnectionPool()
        {
            this.Dispose(false);
        }

        private HostClients HostClientsFromHostName(string host)
        {
            ConnectionPool.HostClients hostClients;
            if (!this.clientsByHost.TryGetValue(host, out hostClients))
            {
                hostClients = new ConnectionPool.HostClients(this, host);
                this.clientsByHost.Add(host, hostClients);
            }
            return hostClients;
        }
        private void UpdateHostClients(IEnumerable<string> hosts)
        {
            HashSet<string> hashSet = new HashSet<string>(this.clientsByHost.Keys);
            foreach (string current in hosts)
            {
                if (!this.clientsByHost.ContainsKey(current))
                {
                    HostClients hostClients = new HostClients(this, current);
                    clientsByHost.Add(current, hostClients);
                }
                hashSet.Remove(current);
            }
            foreach (string current in hashSet)
            {
                HostClients hostClients;
                if (clientsByHost.TryGetValue(current, out hostClients))
                {
                    hostClients.CloseAllIdleClients();
                    clientsByHost.Remove(current);
                }
            }
        }
        private void TryUpdateNodes()
        {
            try
            {
                this.nodeEnumerator.UpdateNodes();
            }
            catch (Exception ex)
            {
            }
        }

        public void UpdateNodes(IEnumerable<string> nodes)
        {
            object borrowReturnLock;
            Monitor.Enter(borrowReturnLock = this.borrowReturnLock);
            try
            {
                this.UpdateHostClients(nodes);
            }
            finally
            {
                Monitor.Exit(borrowReturnLock);
            }
        }
        public Apache.Cassandra.Cassandra.Client Borrow()
        {
            if (this.disposed)
            {
                throw new ObjectDisposedException("ConnectionPool is already disposed.");
            }
            DateTime utcNow = DateTime.UtcNow;
            while (true)
            {
                object borrowReturnLock;
                Monitor.Enter(borrowReturnLock = this.borrowReturnLock);
                try
                {
                    IEnumerable<string> rghost = this.nodeEnumerator.NextNodes();
                    Apache.Cassandra.Cassandra.Client client = this.BorrowCore(rghost, this.MinPoolCountPerHost);
                    if (client != null)
                    {
                        Apache.Cassandra.Cassandra.Client result = client;
                        return result;
                    }
                    List<ConnectionPool.HostClients> rgrgcli = (
                        from cli in this.clientsByHost.Values
                        orderby cli.Count
                        select cli).ToList<ConnectionPool.HostClients>();
                    client = this.BorrowCore(rgrgcli, this.MaxPoolCountPerHost);
                    if (client != null)
                    {
                        Apache.Cassandra.Cassandra.Client result = client;
                        return result;
                    }
                }
                finally
                {
                    Monitor.Exit(borrowReturnLock);
                }
                TimeSpan ts = DateTime.UtcNow.Subtract(utcNow);
                TimeSpan timeSpan = this.Timeout.Subtract(ts);
                if (timeSpan <= TimeSpan.Zero)
                {
                    break;
                }
                this.autoResetEvent.WaitOne(timeSpan);
            }

            throw new TimeoutException("Timeout! Cassandra connections are used to other thread.");
        }
        public void Return(Apache.Cassandra.Cassandra.Client client)
        {
            if (this.disposed)
            {
                throw new ObjectDisposedException("ConnectionPool is already disposed.");
            }
            TSocket tSocket = client.OutputProtocol.Transport.GetTSocket();
            if (tSocket == null)
            {
                throw new ArgumentException("Error, could not get client socket.");
            }
            string host = tSocket.Host;
            if (!tSocket.IsOpen)
            {
                try
                {
                    tSocket.Open();
                }
                catch
                {
                    tSocket.Close();
                    this.TryUpdateNodes();
                    return;
                }
            }
            object borrowReturnLock;
            Monitor.Enter(borrowReturnLock = this.borrowReturnLock);
            try
            {
                ConnectionPool.HostClients hostClients;
                if (!this.clientsByHost.TryGetValue(host, out hostClients))
                {
                    client.Close();
                }
                else
                {
                    hostClients.IdleClients.Enqueue(client);
                    this.autoResetEvent.Set();
                }
            }
            finally
            {
                Monitor.Exit(borrowReturnLock);
            }
        }
        public void ErrorReturn(Apache.Cassandra.Cassandra.Client client, Exception err)
        {
            if (this.disposed)
            {
                throw new ObjectDisposedException("ConnectionPool is already disposed.");
            }
            TSocket tSocket = client.OutputProtocol.Transport.GetTSocket();
            if (tSocket == null)
            {
                throw new ArgumentException("Error, could not get client socket.");
            }
            string host = tSocket.Host;
            client.Close();
            ConnectionPool.HostClients hostClients;
            if (this.clientsByHost.TryGetValue(host, out hostClients))
            {
                hostClients.MarkError(err);
                this.TryUpdateNodes();
            }
        }
        public void Close()
        {
            Dispose();
        }
        public void Dispose()
        {
            Dispose(true);
            GC.SuppressFinalize(this);
        }

        private Apache.Cassandra.Cassandra.Client BorrowCore(IEnumerable<string> rghost, int maxConnections)
        {
            List<ConnectionPool.HostClients> rgrgcli = rghost.Select(new Func<string, ConnectionPool.HostClients>(this.HostClientsFromHostName)).ToList<ConnectionPool.HostClients>();
            return this.BorrowCore(rgrgcli, maxConnections);
        }
        private Apache.Cassandra.Cassandra.Client BorrowCore(ICollection<ConnectionPool.HostClients> rgrgcli, int maxConnections)
        {
            bool fIgnoreRecentError = rgrgcli.All((ConnectionPool.HostClients rgcli) => rgcli.IsLastErrorRecent);
            int num = 0;
            Apache.Cassandra.Cassandra.Client result;
            foreach (ConnectionPool.HostClients current in rgrgcli)
            {
                try
                {
                    Apache.Cassandra.Cassandra.Client client = current.Borrow(maxConnections, fIgnoreRecentError);
                    if (client != null)
                    {
                        result = client;
                        return result;
                    }
                }
                catch (Exception)
                {
                    num++;
                    if (num == rgrgcli.Count)
                    {
                        throw;
                    }
                }
            }
            result = null;
            return result;
        }
        private void CloseAllClient()
        {
            Dictionary<string, ConnectionPool.HostClients> clientsByHost = this.clientsByHost;
            if (clientsByHost != null)
            {
                foreach (ConnectionPool.HostClients current in clientsByHost.Values)
                {
                    current.CloseAllClients();
                }
            }
        }
        private void Dispose(bool disposing)
        {
            if (!this.disposed)
            {
                if (disposing)
                {
                    if (this.autoResetEvent != null)
                    {
                        this.autoResetEvent.Close();
                    }
                    this.CloseAllClient();
                }
                this.autoResetEvent = null;
                this.clientsByHost = null;
                this.disposed = true;
            }
        }

        private static void OnConnected(HostClients hc, Stopwatch sw)
        {
            TimeSpan elapsed = sw.Elapsed;
            EventHandler<CassandraConnectionEventArgs> connected = Connected;
            if (connected != null)
            {
                CassandraConnectionEventArgs e = new CassandraConnectionEventArgs(hc.Host, elapsed, null);
                connected(hc.ConnectionPool, e);
            }
        }
        private static void OnConnectFailed(HostClients hc, Stopwatch sw, Exception err)
        {
            EventHandler<CassandraConnectionEventArgs> connectionFailed = ConnectionFailed;
            if (connectionFailed != null)
            {
                TimeSpan elapsed = sw.Elapsed;
                CassandraConnectionEventArgs e = new CassandraConnectionEventArgs(hc.Host, elapsed, err);
                connectionFailed(hc.ConnectionPool, e);
            }
        }
    }
}
